;;; Ported from Scheme 48 1.9.  See file COPYING for notices and license.
;;;
;;; Port Author: Andrew Whatson
;;;
;;; Original Authors: Richard Kelsey, Mike Sperber
;;;
;;;   scheme48-1.9.2/ps-compiler/node/node-util.scm
;;;
;;; This file contains miscellaneous utilities for accessing and modifying the
;;; node tree.
;;;
;;; Get the root of the tree containing node.

(define-module (ps-compiler node node-util)
  #:use-module (prescheme scheme48)
  #:use-module (ps-compiler node arch)
  #:use-module (ps-compiler node let-nodes)
  #:use-module (ps-compiler node node)
  #:use-module (ps-compiler node primop)
  #:use-module (ps-compiler param)
  #:use-module (ps-compiler util util)
  #:export (node-base containing-procedure
            trivial? nontrivial?
            nontrivial-ancestor
            calls-this-primop?
            bound-to-variable
            walk-refs-safely
            small-node?
            side-effects?
            called-node? called-node
            called-lambda
            get-lambda-value
            ;;set-reference?

            attach-call-args remove-call-args replace-call-args
            remove-null-arguments
            shorten-call-args insert-call-arg remove-call-arg
            append-call-arg

            remove-body
            attach-two-call-args
            attach-three-call-args
            attach-four-call-args
            attach-five-call-args

            remove-lambda-variable remove-variable remove-unused-variables

            substitute substitute-vars-in-node-tree
            replace-call-with-value

            copy-node-tree

            mark-ancestors marked-ancestor? unmarked-ancestor?
            node-ancestor? marked-ancestor least-common-ancestor
            proc-ancestor

            hoistable-node?

            find-scoping

            no-free-references?

            find-calls

            node-type

            the-undefined-value
            undefined-value?
            undefined-value-node?
            make-undefined-literal))

(define (node-base node)
  (do ((p node (node-parent p)))
      ((not (node? (node-parent p)))
       p)))

;; Find the procedure node that contains NODE.  Go up one parent at a time
;; until a lambda node is found, then go up two at a time, skipping the
;; intervening call nodes.

(define (containing-procedure node)
  (do ((node (node-parent node) (node-parent node)))
      ((lambda-node? node)
       (do ((node node (node-parent (node-parent node))))
           ((proc-lambda? node) node)))))

;; Trivial calls are those whose parents are call nodes.

(define (trivial? call)
  (call-node? (node-parent call)))

(define (nontrivial? call)
  (lambda-node? (node-parent call)))

(define (nontrivial-ancestor call)
  (let loop ((call call))
    (if (or (not (node? (node-parent call)))
            (nontrivial? call))
        call
        (loop (node-parent call)))))

(define (calls-this-primop? call id)
  (eq? id (primop-id (call-primop call))))

;; Return the variable to which a value is bound by LET or LETREC.

(define (bound-to-variable node)
  (let ((parent (node-parent node)))
    (case (primop-id (call-primop parent))
      ((let)
       (if (n= 0 (node-index node))
           (list-ref (lambda-variables (call-arg parent 0))
                     (- (node-index node) 1))
           #f))
      ((letrec2)
       (if (< 1 (node-index node))
           (list-ref (lambda-variables
                      (variable-binder
                       (reference-variable (call-arg parent 1))))
                     (- (node-index node) 1))
           #f))
      (else #f))))

;; Return a list of all the reference to lambda-node L's value that call it.
;; If not all can be identified then #F is returned.

(define (find-calls l)
  (let ((refs (cond ((bound-to-variable l)
                     => variable-refs)
                    ((called-node? l)
                     (list l))
                    (else
                     #f))))
    (cond ((and refs (every? called-node? refs))
           refs)
          ((calls-known? l)
           (bug "cannot find calls for known lambda ~S" l))
          (else #f))))

;; Walk (or map) a tree-modifying procedure down a variable's references.

(define (walk-refs-safely proc var)
  (for-each proc (copy-list (variable-refs var))))

;; Return #t if the total primop-cost of NODE is less than SIZE.

(define (small-node? node size)
  (let label ((call (lambda-body node)))
    (set! size (- size (primop-cost call)))
    (if (>= size 0)
        (walk-vector (lambda (n)
                       (cond ((lambda-node? n)
                              (label (lambda-body n)))
                             ((call-node? n)
                              (label n))))
                     (call-args call))))
  (>= size 0))

;; True if executing NODE involves side-effects.

(define (side-effects? node . permissible)
  (let ((permissible (cons #f permissible)))
    (let label ((node node))
      (cond ((not (call-node? node))
             #f)
            ((and (= 0 (call-exits node))
                  (memq (primop-side-effects (call-primop node))
                        permissible))
             (let loop ((i (- (call-arg-count node) 1)))
               (cond ((< i 0) #f)
                     ((label (call-arg node i)) #t)
                     (else (loop (- i 1))))))
            (else
             #t)))))

;; A conservative check - is there only one SET-CONTENTS call for the owner and
;; are all calls between CALL and the LETREC call that binds the owner calls to
;; SET-CONTENTS?

;;(define (single-letrec-set? call)
;;  (let ((owner (call-arg call set/owner)))
;;    (and (reference-node? owner)
;;        (every? (lambda (ref)
;;                  (or (eq? (node-parent ref) call)
;;                      (not (set-reference? ref))))
;;                (variable-refs (reference-variable owner))))))

;;(define (set-reference? node)
;;  (and (eq? 'set-contents
;;           (primop-id (call-primop (node-parent node))))
;;       (= (node-index node) set/owner)))

;;-------------------------------------------------------------------------------

(define the-undefined-value (list '*undefined-value*))

(define (undefined-value? x)
  (eq? x the-undefined-value))

(define (undefined-value-node? x)
  (and (literal-node? x)
       (undefined-value? (literal-value x))))

(define (make-undefined-literal)
  (make-literal-node the-undefined-value #f))

;;-------------------------------------------------------------------------------
;; Finding the lambda node called by CALL, JUMP, or RETURN

(define (called-node? node)
  (and (node? (node-parent node))
       (eq? node (called-node (node-parent node)))))

(define (called-node call)
  (cond ((and (primop-procedure? (call-primop call))
              (primop-call-index (call-primop call)))
         => (lambda (i)
              (call-arg call i)))
        (else '#f)))

(define (called-lambda call)
  (get-lambda-value (call-arg call (primop-call-index (call-primop call)))))

(define (get-lambda-value value)
  (cond ((lambda-node? value)
         value)
        ((reference-node? value)
         (get-variable-lambda (reference-variable value)))
        (else
         (error "peculiar procedure in ~S" value))))

(define (get-variable-lambda variable)
  (if (global-variable? variable)
      (or (variable-known-lambda variable)
          (error "peculiar procedure variable ~S" variable))
      (let* ((binder (variable-binder variable))
             (index (node-index binder))
             (call (node-parent binder))
             (lose (lambda ()
                     (error "peculiar procedure variable ~S" variable))))
        (case (primop-id (call-primop call))
          ((let)
           (if (= 0 index)
               (get-lambda-value (call-arg call (+ 1 (variable-index variable))))
               (lose)))
          ((letrec1)
           (if (= 0 index)
               (get-letrec-variable-lambda variable)
               (lose)))
          ((call)
           (if (and (= 1 index)
                    (= 0 (variable-index variable))) ;; var is a continuation var
               (get-lambda-value (call-arg call 0))
               (lose)))
          (else
           (lose))))))

;; Some of the checking can be removed once I know the LETREC code works.

(define (get-letrec-variable-lambda variable)
  (let* ((binder (variable-binder variable))
         (call (lambda-body binder)))
    (if (and (eq? 'letrec2 (primop-id (call-primop call)))
             (reference-node? (call-arg call 1))
             (eq? (car (lambda-variables binder))
                  (reference-variable (call-arg call 1))))
        (call-arg call (+ 1 (variable-index variable)))
        (error "LETREC is incorrectly organized ~S" (node-parent binder)))))

;;(define (get-cell-variable-lambda variable)
;;  (let ((ref (first set-reference? (variable-refs variable))))
;;    (if (and ref
;;            (eq? 'letrec
;;                 (literal-value (call-arg (node-parent ref) set/type))))
;;       (get-lambda-value (call-arg (node-parent ref) set/value))
;;       (error "peculiar lambda cell ~S" variable))))

;;-------------------------------------------------------------------------------
;; Attaching and detaching arguments to calls

;; Make ARGS the arguments of call node PARENT.  ARGS may contain #f.

(define (attach-call-args parent args)
  (let ((len (call-arg-count parent)))
    (let loop ((args args) (i 0))
      (cond ((null? args)
             (if (< i (- len 1))
                 (bug '"too few arguments added to node ~S" parent))
             (values))
            ((>= i len)
             (bug '"too many arguments added to node ~S" parent))
            (else
             (if (car args)
                 (attach parent i (car args)))
             (loop (cdr args) (+ 1 i)))))))

;; Remove all of the arguments of NODE.

(define (remove-call-args node)
  (let ((len (call-arg-count node)))
    (do ((i 1 (+ i 1)))
        ((>= i len))
      (if (not (empty? (call-arg node i)))
          (erase (detach (call-arg node i)))))
    (values)))

;; Replace the arguments of call node NODE with NEW-ARGS.

(define (replace-call-args node new-args)
  (let ((len (length new-args)))
    (remove-call-args node)
    (if (n= len (call-arg-count node))
        (let ((new (make-vector len empty))
              (old (call-args node)))
          (set-call-args! node new)))
    (attach-call-args node new-args)))

;; Remove all arguments to CALL that are EMPTY?.  COUNT is the number of
;; non-EMPTY? arguments.

(define (remove-null-arguments call count)
  (let ((old (call-args call))
        (new (make-vector count empty)))
    (let loop ((i 0) (j 0))
      (cond ((>= j count)
             (values))
            ((not (empty? (vector-ref old i)))
             (set-node-index! (vector-ref old i) j)
             (vector-set! new j (vector-ref old i))
             (loop (+ i 1) (+ j 1)))
            (else
             (loop (+ i 1) j))))
    (set-call-args! call new)
    (values)))

;; Remove all but the first COUNT arguments from CALL.

(define (shorten-call-args call count)
  (let ((old (call-args call))
        (new (make-vector count empty)))
    (vector-replace new old count)
    (do ((i (+ count 1) (+ i 1)))
        ((>= i (vector-length old)))
      (erase (vector-ref old i)))
    (set-call-args! call new)
    (values)))

;; Insert ARG as the INDEXth argument to CALL.

(define (insert-call-arg call index arg)
  (let* ((old (call-args call))
         (len (vector-length old))
         (new (make-vector (+ 1 len) empty)))
    (vector-replace new old index)
    (do ((i index (+ i 1)))
        ((>= i len))
      (vector-set! new (+ i 1) (vector-ref old i))
      (set-node-index! (vector-ref old i) (+ i 1)))
    (set-call-args! call new)
    (attach call index arg)
    (values)))

;; Remove the INDEXth argument to CALL.

(define (remove-call-arg call index)
  (let* ((old (call-args call))
         (len (- (vector-length old) 1))
         (new (make-vector len)))
    (vector-replace new old index)
    (if (node? (vector-ref old index))
        (erase (detach (vector-ref old index))))
    (do ((i index (+ i 1)))
        ((>= i len))
      (vector-set! new i (vector-ref old (+ i 1)))
      (set-node-index! (vector-ref new i) i))
    (set-call-args! call new)
    (if (< index (call-exits call))
        (set-call-exits! call (- (call-exits call) 1)))
    (values)))

;; Add ARG to the end of CALL's arguments.

(define (append-call-arg call arg)
  (insert-call-arg call (call-arg-count call) arg))

;; Replace CALL with the body of its continuation.

(define (remove-body call)
  (if (n= 1 (call-exits call))
      (bug "removing a call with ~D exits" (call-exits call))
      (replace-body call (detach-body (lambda-body (call-arg call 0))))))

;; Avoiding N-Ary Procedures
;; These are used in the expansion of the LET-NODES macro.

(define (attach-two-call-args node a0 a1)
  (attach node 0 a0)
  (attach node 1 a1))

(define (attach-three-call-args node a0 a1 a2)
  (attach node 0 a0)
  (attach node 1 a1)
  (attach node 2 a2))

(define (attach-four-call-args node a0 a1 a2 a3)
  (attach node 0 a0)
  (attach node 1 a1)
  (attach node 2 a2)
  (attach node 3 a3))

(define (attach-five-call-args node a0 a1 a2 a3 a4)
  (attach node 0 a0)
  (attach node 1 a1)
  (attach node 2 a2)
  (attach node 3 a3)
  (attach node 4 a4))

;;-------------------------------------------------------------------------------
;; Changing lambda-nodes' variable lists

(define (remove-lambda-variable l-node index)
  (remove-variable l-node (list-ref (lambda-variables l-node) index)))

(define (remove-variable l-node var)
  (if (used? var)
      (bug '"cannot remove referenced variable ~s" var))
  (erase-variable var)
  (let ((vars (lambda-variables l-node)))
    (if (eq? (car vars) var)
        (set-lambda-variables! l-node (cdr vars))
        (do ((vars vars (cdr vars)))
            ((eq? (cadr vars) var)
             (set-cdr! vars (cddr vars)))))))

;; Remove all of L-NODES' unused variables.

(define (remove-unused-variables l-node)
  (set-lambda-variables! l-node
                         (filter! (lambda (v)
                                    (cond ((used? v)
                                           #t)
                                          (else
                                           (erase-variable v)
                                           #f)))
                                  (lambda-variables l-node))))

;;------------------------------------------------------------------------------
;; Substituting Values For Variables

;; Substitute VAL for VAR.  If DETACH? is true then VAL should be detached
;; and so can be used instead of a copy for the first substitution.
;;
;; If VAL is a reference to a variable named V, it was probably introduced by
;; the CPS conversion code.  In that case, the variable is renamed with the
;; name of VAR.  This helps considerably when debugging the compiler.

(define (substitute var val detach?)
  (if (and (reference-node? val)
           (eq? 'v (variable-name (reference-variable val)))
           (not (global-variable? (reference-variable val))))
      (set-variable-name! (reference-variable val)
           (variable-name var)))
  (let ((refs (variable-refs var)))
    (set-variable-refs! var '())
    (cond ((not (null? refs))
           (for-each (lambda (ref)
                       (replace ref (copy-node-tree val)))
                     (if detach? (cdr refs) refs))
           (if detach? (replace (car refs) (detach val))))
          (detach?
           (erase (detach val))))))

;; Walk the tree NODE replacing references to variables in OLD-VARS with
;; the corresponding variables in NEW-VARS.  Uses VARIABLE-FLAG to mark
;; the variables being replaced.

(define (substitute-vars-in-node-tree node old-vars new-vars)
  (for-each (lambda (old new)
              (set-variable-flag! old new))
            old-vars
            new-vars)
  (let tree-walk ((node node))
    (cond ((lambda-node? node)
           (walk-vector tree-walk (call-args (lambda-body node))))
          ((call-node? node)
           (walk-vector tree-walk (call-args node)))
          ((and (reference-node? node)
                (variable-flag (reference-variable node)))
           => (lambda (new)
                (replace node (make-reference-node new))))))
  (for-each (lambda (old)
              (set-variable-flag! old #f))
            old-vars))

;; Replaces the call node CALL with VALUE.
;; (<proc> <exit> . <args>) => (<exit> <value>)

(define (replace-call-with-value call value)
  (cond ((n= 1 (call-exits call))
         (bug '"can only substitute for call with one exit ~s" call))
        (else
         (let ((cont (detach (call-arg call 0))))
           (set-call-exits! call 0)
           (replace-call-args call (if value (list cont value) (list cont)))
           (set-call-primop! call (get-primop (enum primop-enum let)))))))

;;------------------------------------------------------------------------------
;; Copying Node Trees

;; Copy the node-tree NODE.  This dispatches on the type of NODE.
;; Variables which have been copied have the copy in the node-flag field.

(define (copy-node-tree node)
  (let ((new (cond ((lambda-node? node)
                    (copy-lambda node))
                   ((reference-node? node)
                    (let ((var (reference-variable node)))
                      (cond ((and (variable-binder var)
                                  (variable-flag var))
                             => make-reference-node)
                            (else
                             (make-reference-node var)))))
                   ((call-node? node)
                    (copy-call node))
                   ((literal-node? node)
                    (copy-literal-node node)))))
    new))

;; Copy a lambda node and its variables.  The variables' copies are put in
;; their VARIABLE-FLAG while the lambda's body is being copied.

(define (copy-lambda node)
  (let* ((vars (map (lambda (var)
                      (if var
                          (let ((new (copy-variable var)))
                            (set-variable-flag! var new)
                            new)
                          #f))
                    (lambda-variables node)))
         (new-node (make-lambda-node (lambda-name node)
                                     (lambda-type node)
                                     vars)))
    (attach-body new-node (copy-call (lambda-body node)))
    (set-lambda-protocol! new-node (lambda-protocol node))
    (set-lambda-source! new-node (lambda-source node))
    (for-each (lambda (var)
                (if var (set-variable-flag! var #f)))
              (lambda-variables node))
    new-node))

(define (copy-call node)
  (let ((new-node (make-call-node (call-primop node)
                                  (call-arg-count node)
                                  (call-exits node))))
    (do ((i 0 (+ i 1)))
        ((>= i (call-arg-count node)))
      (attach new-node i (copy-node-tree (call-arg node i))))
    (set-call-source! new-node (call-source node))
    new-node))

;;------------------------------------------------------------------------------
;; Checking the scoping of identifers

;; Mark all ancestors of N with FLAG

(define (mark-ancestors n flag)
  (do ((n n (node-parent n)))
      ((not (node? n)) (values))
    (set-node-flag! n flag)))

;; Does N have an ancestor with a non-#f flag?

(define (marked-ancestor? n)
  (do ((n n (node-parent n)))
      ((or (not (node? n))
           (node-flag n))
       (node? n))))

;; Does N have an ancestor with a #f flag?

(define (unmarked-ancestor? n)
  (do ((n n (node-parent n)))
      ((or (not (node? n))
           (not (node-flag n)))
       (node? n))))

;; Is ANC? an ancestor of NODE?

(define (node-ancestor? anc? node)
  (set-node-flag! anc? #t)
  (let ((okay? (marked-ancestor? node)))
    (set-node-flag! anc? #f)
    okay?))

;; Find the lowest ancestor of N that has a non-#f flag

(define (marked-ancestor n)
  (do ((n n (node-parent n)))
      ((or (not (node? n))
           (node-flag n))
       (if (node? n) n #f))))

;; Mark the ancestors of START with #f, stopping when END is reached

(define (unmark-ancestors-to start end)
  (do ((node start (node-parent node)))
      ((eq? node end))
    (set-node-flag! node #f)))

;; Return the lowest node that is above all NODES

(define (least-common-ancestor nodes)
  (mark-ancestors (car nodes) #t)
  (let loop ((nodes (cdr nodes)) (top (car nodes)))
    (cond ((null? nodes)
           (mark-ancestors top #f)
           top)
          (else
           (let ((new (marked-ancestor (car nodes))))
             (unmark-ancestors-to top new)
             (loop (cdr nodes) new))))))

;; Can TO be moved to FROM without taking variables out of scope.
;; This first marks all of the ancestors of FROM, and then unmarks all of the
;; ancestors of TO.  The net result is to mark every node that is above FROM but
;; not above TO.  Then if any reference-node below FROM references a variable
;; with a marked binder, that node, and thus FROM itself, cannot legally be
;; moved to TO.

;; This is not currently used anywhere, and it doesn't know about trivial
;; calls.

(define (hoistable-node? from to)
  (let ((from (if (call-node? from)
                  (node-parent (nontrivial-ancestor from))
                  from)))
    (mark-ancestors (node-parent from) #t)
    (mark-ancestors to #f)
    (let ((okay? (let label ((n from))
                   (cond ((lambda-node? n)
                          (let* ((vec (call-args (lambda-body n)))
                                 (c (vector-length vec)))
                            (let loop ((i 0))
                              (cond ((>= i c) #t)
                                    ((label (vector-ref vec i))
                                     (loop (+ i 1)))
                                    (else #f)))))
                         ((reference-node? n)
                          (let ((b (variable-binder (reference-variable n))))
                            (or (not b) (not (node-flag b)))))
                         (else #t)))))
      (mark-ancestors (node-parent from) #f)
      okay?)))

;; Mark all of the lambda nodes which bind variables referenced below NODE.

(define (mark-binders node)
  (let label ((n node))
    (cond ((lambda-node? n)
           (walk-vector label (call-args (lambda-body n))))
          ((reference-node? n)
           (let ((b (variable-binder (reference-variable n))))
             (if b (set-node-flag! b #f))))))
  (values))


;;------------------------------------------------------------------------------
;; For each lambda-node L this sets (PARENT L) to be the enclosing PROC node
;; of L and, if L is a PROC node, sets (KIDS L) to be the lambda nodes it
;; encloses.

(define (find-scoping lambdas parent set-parent! kids set-kids!)
  (receive (procs others)
      (partition-list proc-lambda? lambdas)
    (for-each (lambda (l)
                (set-parent! l #f)
                (set-kids!   l '()))
              procs)
    (for-each (lambda (l)
                (set-parent! l #f))
              others)
    (letrec ((set-lambda-parent!
              (lambda (l)
                (cond ((parent l)
                       => identity)
                      ((proc-ancestor l)
                       => (lambda (p)
                            (let ((p (if (proc-lambda? p)
                                         p
                                         (set-lambda-parent! p))))
                              (set-kids! p (cons l (kids p)))
                              (set-parent! l p)
                              p)))
                      (else #f)))))
      (for-each set-lambda-parent! lambdas))
    (values procs others)))

(define (proc-ancestor node)
  (let ((p (node-parent node)))
    (if (not (node? p))
        #f
        (let ((node (do ((p p (node-parent p)))
                        ((lambda-node? p)
                         p))))
          (do ((node node (node-parent (node-parent node))))
              ((proc-lambda? node)
               node))))))

(define (no-free-references? node)
  (if (call-node? node)
      (error "NO-FREE-REFERENCES only works on value nodes: ~S" node))
  (let label ((node node))
    (cond ((reference-node? node)
           (let ((b (variable-binder (reference-variable node))))
             (or (not b)
                 (node-flag b))))
          ((lambda-node? node)
           (set-node-flag! node #t)
           (let ((res (label (lambda-body node))))
             (set-node-flag! node #f)
             res))
          ((call-node? node)
           (let ((vec (call-args node)))
             (let loop ((i (- (vector-length vec) 1)))
               (cond ((< i 0) #t)
                     ((not (label (vector-ref vec i))) #f)
                     (else (loop (- i 1)))))))
          (else #t))))

(define (node-type node)
  (cond ((literal-node? node)
         (literal-type node))
        ((reference-node? node)
         (variable-type (reference-variable node)))
        ((lambda-node? node)
         (lambda-node-type node))
        ((and (call-node? node)
              (primop-trivial? (call-primop node)))
         (trivial-call-return-type node))
        (else
         (error "node ~S does not represent a value" node))))

;;----------------------------------------------------------------
;; Debugging utilities

(define (show-simplified node)
  (let loop ((n node) (r '()))
    (if (node? n)
        (loop (node-parent n) (cons (node-simplified? n) r))
        (reverse r))))

(define (show-flag node)
  (let loop ((n node) (r '()))
    (if (node? n)
        (loop (node-parent n) (cons (node-flag n) r))
        (reverse r))))

(define (reset-simplified node)
  (let loop ((n node))
    (cond ((node? n)
           (set-node-simplified?! n #f)
           (loop (node-parent n))))))
